from easycore.common.config import CfgNode

def get_dataset_info_cfg(path):
    cfg = CfgNode.open(path)
    return cfg

def add_params_to_dataset_dicts(dataset_dicts, params:dict):
    for i in range(len(dataset_dicts)):
        for key, value in params.items():
            dataset_dicts[i]['params'][key] = value
    return dataset_dicts

def setup_dataset_info(dataset_dicts, cfg):
    for i in range(len(dataset_dicts)):
        dataset_dicts[i]['preprocess'] = cfg.preprocess
        for key, value in cfg.params.items():
            dataset_dicts[i]['params'][key] = value
    return dataset_dicts

def convert_to_one_category(dataset_dicts, category_id=0):
    for i in range(len(dataset_dicts)):
        for j in range(len(dataset_dicts[i]['annotations'])):
            dataset_dicts[i]['annotations'][j]['category_id'] = category_id
    return dataset_dicts

def map_category(dataset_dicts, map_dict):
    for i in range(len(dataset_dicts)):
        for j in range(len(dataset_dicts[i]['annotations'])):
            dataset_dicts[i]['annotations'][j]['category_id'] = map_dict[dataset_dicts[i]['annotations'][j]['category_id']]
    return dataset_dicts

def get_plugin_annotations(annotations_file):
    annotations_dict = {}
    with open(annotations_file, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip().split('\t')
            filename = line[0]
            num_objs = int(line[1])
            annotations = []
            for obj_text in line[2:]:
                obj = obj_text.split(' ')
                object_dict = {
                    'bbox': [float(i) for i in obj[:4]],
                    'category_id': int(obj[4]),
                }
                annotations.append(object_dict)
            annotations_dict[filename] = annotations
    return annotations_dict
                
